#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : RAG-demo
# @File    : milvus_util.py
# @IDE     : PyCharm
# @Author  :ZH
# @Time    : 2025/6/27 17:13
from typing import Union, List

from pymilvus import MilvusClient, FieldSchema, DataType, CollectionSchema
from pymilvus.client.types import LoadState

from util.logger_setup_util import LoggerSetup

logger = LoggerSetup(__name__).get_logger()


class MethodWrapper:
    def __init__(self, func):
        self.func = func

    def __get__(self, instance, owner):
        def wrapper(*args, **kwargs):
            instance.supplement()
            try:
                return self.func(instance, *args, **kwargs)
            finally:
                instance.close_client()

        return wrapper


class KnowledgeMilvus:
    __doc__ = '知识库Milvus 工具类'

    def __init__(self, uri, user='', password='', db_name='', token=''):
        self.uri = uri
        self.user = user
        self.password = password
        self.db_name = db_name
        self.token = token
        self.client = None

    def create_client(self):
        logger.info('创建向量库链接')
        self.client = MilvusClient(
            uri=self.uri,
            user=self.user,
            password=self.password,
            db_name=self.db_name,
            token=self.token
        )

    def is_client(self):
        if self.client:
            return True
        else:
            return False

    def close_client(self):
        logger.info('关闭向量库链接')
        try:
            if self.client:
                self.client.close()
            self.client = None
        except Exception as e:
            logger.error('close client error', e)

    def supplement(self):
        if not self.client:
            self.create_client()

    @MethodWrapper
    def create_collection(self, collection_name, dim=768, description=''):

        # 判断是否存在，如果存在就不进行创建
        if self.client.has_collection(collection_name=collection_name):
            return
        # 设置id字段并且设置id自动产生
        document_id = FieldSchema(name='document_id', dtype=DataType.INT64, is_primary=True, auto_id=True)
        # 设置向量字段
        vector = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim)
        # 设置文件名称
        file_name = FieldSchema(name="file_name", dtype=DataType.VARCHAR, max_length=256)
        # 文段内容
        document = FieldSchema(name="document", dtype=DataType.VARCHAR, max_length=8192)
        # 文件hash
        file_hase = FieldSchema(name="file_hase", dtype=DataType.VARCHAR, max_length=256)
        # 知识库id
        knowledge_id = FieldSchema(name='knowledge_id', dtype=DataType.INT64)
        # 设置向量库索引属性
        index_type = 'IVF_FLAT'
        # index_type = 'FLAT'
        metric_type = 'IP'
        params = {"nlist": 1024}
        schema = CollectionSchema(
            fields=[document_id, vector, file_name, document, file_hase, knowledge_id],
            description=description)
        index_param = self.client.prepare_index_params()
        # 为向量字段添加索引
        index_param.add_index(field_name='vector', index_type=index_type, params=params, metric_type=metric_type)
        # 标量字段添加索引
        index_param.add_index(field_name='knowledge_id', index_type='STL_SORT')
        self.client.create_collection(
            collection_name=collection_name, dimension=dim, schema=schema,
        )
        self.client.create_index(collection_name=collection_name, index_params=index_param)

    @MethodWrapper
    def load_collection(self, collection_name):
        """
        加载向量库
        :param collection_name: 向量库名称
        :return:
        """

        state = self.client.get_load_state(collection_name).get('state')
        if state == LoadState.Loaded:
            return True
        elif state == LoadState.NotLoad:
            self.client.load_collection(collection_name=collection_name)
            return True
        else:
            return False

    @MethodWrapper
    def create_partitions(self, collection_name, partition_name, is_load: bool = False):
        """
        创建分区
        Args:
            collection_name: 向量库名称
            partition_name: 分区名称
            is_load:

        Returns:

        """

        if self.client.has_partition(collection_name=collection_name, partition_name=partition_name):
            return
        self.client.create_partition(collection_name=collection_name, partition_name=partition_name)
        if is_load:
            self.load_partitions(collection_name=collection_name, partition_name=partition_name)

    @MethodWrapper
    def load_partitions(self, collection_name, partition_name):
        """
        加载分区
        Args:
            collection_name:
            partition_name:

        Returns:

        """

        state = self.client.get_load_state(collection_name, partition_name=partition_name).get('state')
        if state == LoadState.Loaded:

            return True
        elif state == LoadState.NotLoad:
            self.client.load_partitions(partition_names=partition_name, collection_name=collection_name)

            return True
        else:

            return False

    @MethodWrapper
    def insert(self, collection_name, data: Union[dict, List[dict]], partition_name: str = None):
        """
        插入向量
        Args:
            collection_name: 向量库名称
            data: 数据
            partition_name: 分区名称

        Returns:

        """
        if isinstance(data, list):
            data_list = [data[i:i + 100] for i in range(0, len(data), 100)]
            # 数据为list的时候 我们需要进行分批录入
            for index, data in enumerate(data_list):
                try:

                    logger.info(f'录入{index + 1}批')
                    self.client.insert(collection_name=collection_name, data=data, partition_name=partition_name)
                except Exception as e:
                    logger.error(f'录入向量库 第{index + 1}批 error', e)
        else:
            self.client.insert(collection_name=collection_name, data=data, partition_name=partition_name)

        return True

    @MethodWrapper
    def search(self, collection_name, data: Union[List[list], list], output_fields: list, filter_str: str = '',
               partition_names: List[str] = None, limit: int = 10, offset: int = 0):
        """
        搜索向量
        Args:
            collection_name: 向量库名称
            data: 向量
            output_fields: 需要输出的字段
            filter_str: 字符查询条件按
            partition_names: partition_name列表
            limit: 返回条数
            offset: 开始的条数

        Returns:

        """

        search_params = {"metric_type": "IP", "params": {"nprobe": 10}}
        res = self.client.search(
            collection_name=collection_name, data=data, limit=limit, output_fields=output_fields,
            search_params=search_params, partition_names=partition_names,
            filter=filter_str, offset=offset, anns_field="vector"
        )

        return res

    @MethodWrapper
    def query(self, collection_name, filter_str: str = '', output_fields: list = None, limit: int = None, ids=None,
              partition_names=None, offset: int = None):
        """
        查询向量
        Args:
            collection_name: 向量库名称
            filter_str: 查询字符串
            output_fields: 输出字段限制
            limit: 查询条数
            ids: 限制id
            partition_names: 分区名称
            offset: offset

        Returns:

        """

        res = self.client.query(
            collection_name=collection_name, filter=filter_str, output_fields=output_fields,
            limit=limit, ids=ids, partition_names=partition_names, offset=offset
        )

        return res

    @MethodWrapper
    def delete(self, collection_name, ids: Union[list, str, int] = None, partition_name: str = '',
               filter_str: str = ""):
        """
        向量删除
        Args:
            collection_name:
            ids:
            partition_name:
            filter_str:

        Returns:

        """
        if not ids and not filter_str:
            raise Exception('条件不能为空')

        return self.client.delete(collection_name=collection_name, ids=ids, partition_name=partition_name,
                                  filter=filter_str)

    @MethodWrapper
    def drop_partition(self, collection_name: str, partition_name: str):
        """
        删除分区
        :param collection_name:
        :param partition_name:
        :return:
        """
        if self.client.has_partition(collection_name=collection_name, partition_name=partition_name):
            state = self.client.get_load_state(collection_name, partition_name=partition_name).get('state')
            if state == LoadState.Loaded:
                self.client.release_partitions(collection_name=collection_name, partition_names=[partition_name])
            self.client.drop_partition(collection_name=collection_name, partition_name=partition_name)
